#!/usr/bin/env python3
"""
WebRTC测试服务器
用于测试服务器环境对WebRTC的支持
包含详细的日志记录以帮助定位问题
"""

import asyncio
import json
import asyncio
import json
import logging
import os
import time
import fractions
from datetime import datetime
from aiohttp import web, WSMsgType
import aiohttp_cors
from aiortc import RTCPeerConnection, RTCSessionDescription, VideoStreamTrack, AudioStreamTrack, MediaStreamTrack, RTCConfiguration, RTCIceServer
import cv2
import numpy as np

# 配置详细日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - [%(filename)s:%(lineno)d] - %(message)s',
    handlers=[
        logging.FileHandler('webrtc_test.log'),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

class TestVideoTrack(VideoStreamTrack):
    """
    测试视频轨道 - 生成彩色测试图案
    """
    def __init__(self):
        super().__init__()
        self.counter = 0
        logger.info("TestVideoTrack initialized")

    async def recv(self):
        pts, time_base = await self.next_timestamp()
        
        # 创建测试图案
        frame = np.zeros((480, 640, 3), dtype=np.uint8)
        
        # 添加彩色条纹和计数器
        color = (self.counter % 255, (self.counter * 2) % 255, (self.counter * 3) % 255)
        cv2.rectangle(frame, (0, 0), (640, 480), color, -1)
        
        # 添加文字
        text = f"Frame: {self.counter} Time: {datetime.now().strftime('%H:%M:%S')}"
        cv2.putText(frame, text, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2)
        
        self.counter += 1
        
        # 转换为VideoFrame
        from av import VideoFrame
        av_frame = VideoFrame.from_ndarray(frame, format='bgr24')
        av_frame.pts = pts
        av_frame.time_base = time_base
        
        if self.counter % 30 == 0:  # 每30帧记录一次
            logger.debug(f"Generated video frame {self.counter}")
        
        return av_frame

class TestAudioTrack(AudioStreamTrack):
    """
    测试音频轨道 - 生成正弦波音频
    """
    def __init__(self):
        super().__init__()
        self.sample_rate = 48000
        self.samples = 0
        self._start = None
        logger.info("TestAudioTrack initialized")

    async def recv(self):
        if self._start is None:
            self._start = time.time()
        
        # 计算时间戳
        now = time.time()
        pts = int((now - self._start) * self.sample_rate)
        time_base = fractions.Fraction(1, self.sample_rate)
        
        # 生成1秒的音频数据
        samples_per_frame = int(self.sample_rate * 0.02)  # 20ms
        
        # 生成440Hz正弦波
        t = np.linspace(self.samples / self.sample_rate, 
                       (self.samples + samples_per_frame) / self.sample_rate, 
                       samples_per_frame, False)
        audio_data = np.sin(2 * np.pi * 440 * t) * 0.3  # 440Hz, 30%音量
        
        self.samples += samples_per_frame
        
        # 转换为AudioFrame
        from av import AudioFrame
        av_frame = AudioFrame.from_ndarray(
            audio_data.reshape(1, -1).astype(np.float32), 
            format='flt', 
            layout='mono'
        )
        av_frame.sample_rate = self.sample_rate
        av_frame.pts = pts
        av_frame.time_base = time_base
        
        if self.samples % (self.sample_rate * 2) < samples_per_frame:  # 每2秒记录一次
            logger.debug(f"Generated audio frame, total samples: {self.samples}")
        
        return av_frame

class WebRTCServer:
    def __init__(self, use_stun=True, stun_servers=None):
        self.connections = {}
        self.connection_stats = {}
        self.use_stun = use_stun
        self.stun_servers = stun_servers or [
            "stun:stun.l.google.com:19302",
            "stun:stun1.l.google.com:19302"
        ]
        logger.info(f"WebRTC Server initialized with STUN: {use_stun}")
        if use_stun:
            logger.info(f"STUN servers: {self.stun_servers}")
    
    async def websocket_handler(self, request):
        """WebSocket处理器"""
        ws = web.WebSocketResponse()
        await ws.prepare(request)
        
        client_id = id(ws)
        client_ip = request.remote
        logger.info(f"New WebSocket connection from {client_ip}, client_id: {client_id}")
        
        self.connections[client_id] = {
            'ws': ws,
            'pc': None,
            'ip': client_ip,
            'connect_time': time.time(),
            'last_activity': time.time()
        }
        
        try:
            async for msg in ws:
                if msg.type == WSMsgType.TEXT:
                    await self.handle_message(client_id, json.loads(msg.data))
                elif msg.type == WSMsgType.ERROR:
                    logger.error(f"WebSocket error for client {client_id}: {ws.exception()}")
                    break
        except Exception as e:
            logger.error(f"Error handling WebSocket for client {client_id}: {e}")
        finally:
            await self.cleanup_connection(client_id)
        
        return ws
    
    async def handle_message(self, client_id, message):
        """处理WebSocket消息"""
        try:
            msg_type = message.get('type')
            logger.info(f"Received message from client {client_id}: {msg_type}")
            logger.debug(f"Message content: {message}")
            
            self.connections[client_id]['last_activity'] = time.time()
            
            if msg_type == 'offer':
                await self.handle_offer(client_id, message)
            elif msg_type == 'answer':
                await self.handle_answer(client_id, message)
            elif msg_type == 'ice-candidate':
                await self.handle_ice_candidate(client_id, message)
            elif msg_type == 'ping':
                await self.send_message(client_id, {'type': 'pong', 'timestamp': time.time()})
            else:
                logger.warning(f"Unknown message type: {msg_type}")
                
        except Exception as e:
            logger.error(f"Error handling message from client {client_id}: {e}")
    
    async def handle_offer(self, client_id, message):
        """处理WebRTC offer"""
        try:
            logger.info(f"Processing offer from client {client_id}")
            
            # 创建RTCPeerConnection配置
            if self.use_stun:
                ice_servers = [RTCIceServer(urls=server) for server in self.stun_servers]
                config = RTCConfiguration(iceServers=ice_servers)
                logger.info(f"Client {client_id} using STUN servers: {self.stun_servers}")
            else:
                config = RTCConfiguration(iceServers=[])
                logger.info(f"Client {client_id} not using STUN servers (host-only mode)")
            
            # 创建RTCPeerConnection
            pc = RTCPeerConnection(configuration=config)
            self.connections[client_id]['pc'] = pc
            
            # 设置事件处理器
            @pc.on("connectionstatechange")
            async def on_connectionstatechange():
                logger.info(f"Client {client_id} connection state: {pc.connectionState}")
                if pc.connectionState == "connected":
                    logger.info(f"Client {client_id} successfully connected via WebRTC")
                elif pc.connectionState == "failed":
                    logger.error(f"Client {client_id} WebRTC connection failed")
                elif pc.connectionState == "disconnected":
                    logger.info(f"Client {client_id} WebRTC connection disconnected")
            
            @pc.on("iceconnectionstatechange")
            async def on_iceconnectionstatechange():
                logger.info(f"Client {client_id} ICE connection state: {pc.iceConnectionState}")
                if pc.iceConnectionState == "connected":
                    logger.info(f"Client {client_id} ICE connection established")
                elif pc.iceConnectionState == "failed":
                    logger.error(f"Client {client_id} ICE connection failed")
            
            @pc.on("icegatheringstatechange")
            async def on_icegatheringstatechange():
                logger.info(f"Client {client_id} ICE gathering state: {pc.iceGatheringState}")
            
            @pc.on("track")
            def on_track(track):
                logger.info(f"Client {client_id} received track: {track.kind}")
            
            # 添加测试轨道
            logger.info(f"Adding test video and audio tracks for client {client_id}")
            video_track = TestVideoTrack()
            audio_track = TestAudioTrack()
            
            pc.addTrack(video_track)
            pc.addTrack(audio_track)
            
            # 设置远程描述
            offer_sdp = message['sdp']
            logger.debug(f"Client {client_id} offer SDP: {offer_sdp}")
            
            await pc.setRemoteDescription(RTCSessionDescription(sdp=offer_sdp, type="offer"))
            logger.info(f"Set remote description for client {client_id}")
            
            # 创建answer
            answer = await pc.createAnswer()
            await pc.setLocalDescription(answer)
            logger.info(f"Created answer for client {client_id}")
            logger.debug(f"Client {client_id} answer SDP: {answer.sdp}")
            
            # 发送answer
            await self.send_message(client_id, {
                'type': 'answer',
                'sdp': answer.sdp
            })
            
            logger.info(f"Sent answer to client {client_id}")
            
        except Exception as e:
            logger.error(f"Error handling offer from client {client_id}: {e}")
            await self.send_message(client_id, {
                'type': 'error',
                'message': f"Failed to process offer: {str(e)}"
            })
    
    async def handle_answer(self, client_id, message):
        """处理WebRTC answer"""
        try:
            logger.info(f"Processing answer from client {client_id}")
            pc = self.connections[client_id]['pc']
            if pc:
                await pc.setRemoteDescription(RTCSessionDescription(sdp=message['sdp'], type="answer"))
                logger.info(f"Set remote description (answer) for client {client_id}")
        except Exception as e:
            logger.error(f"Error handling answer from client {client_id}: {e}")
    
    async def handle_ice_candidate(self, client_id, message):
        """处理ICE candidate"""
        try:
            logger.debug(f"Processing ICE candidate from client {client_id}")
            pc = self.connections[client_id]['pc']
            if pc:
                candidate_data = message['candidate']
                logger.debug(f"Client {client_id} ICE candidate: {candidate_data}")
                
                # 如果candidate为空或None，跳过处理
                if not candidate_data or not candidate_data.get('candidate'):
                    logger.debug(f"Skipping empty ICE candidate for client {client_id}")
                    return
                
                # 获取candidate字符串
                candidate_str = candidate_data.get('candidate', '')
                if not candidate_str:
                    logger.debug(f"Empty candidate string for client {client_id}")
                    return
                
                from aiortc import RTCIceCandidate
                
                # 解析candidate字符串，例如: "candidate:1 1 UDP 2130706431 192.168.1.100 54400 typ host"
                parts = candidate_str.split()
                if len(parts) < 8:
                    logger.warning(f"Invalid candidate format for client {client_id}: {candidate_str}")
                    return
                
                # 根据aiortc文档构造RTCIceCandidate
                # class aiortc.RTCIceCandidate(component, foundation, ip, port, priority, protocol, type, ...)
                candidate = RTCIceCandidate(
                    component=int(parts[1]),
                    foundation=parts[0].split(':')[1],
                    ip=parts[4],
                    port=int(parts[5]),
                    priority=int(parts[3]),
                    protocol=parts[2].upper(),  # UDP/TCP
                    type=parts[7],  # host/srflx/relay等
                    sdpMid=candidate_data.get('sdpMid'),
                    sdpMLineIndex=candidate_data.get('sdpMLineIndex')
                )
                
                await pc.addIceCandidate(candidate)
                logger.debug(f"Added ICE candidate for client {client_id}")
        except Exception as e:
            logger.error(f"Error handling ICE candidate from client {client_id}: {e}")
    
    async def send_message(self, client_id, message):
        """发送消息给客户端"""
        try:
            if client_id in self.connections:
                ws = self.connections[client_id]['ws']
                await ws.send_str(json.dumps(message))
                logger.debug(f"Sent message to client {client_id}: {message['type']}")
        except Exception as e:
            logger.error(f"Error sending message to client {client_id}: {e}")
    
    async def cleanup_connection(self, client_id):
        """清理连接"""
        try:
            if client_id in self.connections:
                connection = self.connections[client_id]
                if connection['pc']:
                    await connection['pc'].close()
                    logger.info(f"Closed RTCPeerConnection for client {client_id}")
                
                connect_duration = time.time() - connection['connect_time']
                logger.info(f"Client {client_id} disconnected after {connect_duration:.2f} seconds")
                
                del self.connections[client_id]
        except Exception as e:
            logger.error(f"Error cleaning up connection for client {client_id}: {e}")
    
    async def stats_handler(self, request):
        """统计信息处理器"""
        stats = {
            'active_connections': len(self.connections),
            'connections': []
        }
        
        for client_id, conn in self.connections.items():
            conn_info = {
                'client_id': client_id,
                'ip': conn['ip'],
                'connect_time': conn['connect_time'],
                'duration': time.time() - conn['connect_time'],
                'last_activity': conn['last_activity'],
                'pc_state': conn['pc'].connectionState if conn['pc'] else 'none',
                'ice_state': conn['pc'].iceConnectionState if conn['pc'] else 'none'
            }
            stats['connections'].append(conn_info)
        
        return web.json_response(stats)

async def create_app(use_stun=True, stun_servers=None):
    """创建web应用"""
    logger.info("Creating web application")
    
    app = web.Application()
    server = WebRTCServer(use_stun=use_stun, stun_servers=stun_servers)
    
    # 设置CORS
    cors = aiohttp_cors.setup(app, defaults={
        "*": aiohttp_cors.ResourceOptions(
            allow_credentials=True,
            expose_headers="*",
            allow_headers="*",
            allow_methods="*"
        )
    })
    
    # 路由
    app.router.add_get('/ws', server.websocket_handler)
    app.router.add_get('/stats', server.stats_handler)
    app.router.add_static('/', path='.', name='static')
    
    # 添加CORS到所有路由
    for route in list(app.router.routes()):
        cors.add(route)
    
    logger.info("Web application created successfully")
    return app

async def main():
    """主函数"""
    import argparse
    
    # 解析命令行参数
    parser = argparse.ArgumentParser(description='WebRTC测试服务器')
    parser.add_argument('--no-stun', action='store_true', help='禁用STUN服务器（仅使用host候选）')
    parser.add_argument('--stun-servers', nargs='+', help='自定义STUN服务器列表')
    parser.add_argument('--port', type=int, default=8080, help='服务器端口 (默认: 8080)')
    args = parser.parse_args()
    
    # 配置STUN
    use_stun = not args.no_stun
    stun_servers = args.stun_servers
    
    logger.info("Starting WebRTC test server")
    logger.info(f"Python version: {os.sys.version}")
    logger.info(f"Working directory: {os.getcwd()}")
    logger.info(f"STUN enabled: {use_stun}")
    if use_stun and stun_servers:
        logger.info(f"Custom STUN servers: {stun_servers}")
    
    try:
        app = await create_app(use_stun=use_stun, stun_servers=stun_servers)
        
        # 启动服务器
        runner = web.AppRunner(app)
        await runner.setup()
        
        site = web.TCPSite(runner, '0.0.0.0', args.port)
        await site.start()
        
        logger.info("Server started successfully")
        logger.info(f"Server listening on http://0.0.0.0:{args.port}")
        logger.info(f"WebSocket endpoint: ws://localhost:{args.port}/ws")
        logger.info(f"Stats endpoint: http://localhost:{args.port}/stats")
        logger.info("Press Ctrl+C to stop")
        
        # 保持服务器运行
        try:
            await asyncio.Future()  # 永远等待
        except KeyboardInterrupt:
            logger.info("Received shutdown signal")
        finally:
            await runner.cleanup()
            logger.info("Server stopped")
            
    except Exception as e:
        logger.error(f"Failed to start server: {e}")
        raise

if __name__ == '__main__':
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        logger.info("Server interrupted by user")
    except Exception as e:
        logger.error(f"Server error: {e}")
        raise